#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
预测结果融合机制
实现多种算法预测结果的加权平均融合
"""

import numpy as np
from collections import defaultdict

class PredictionFusion:
    """预测结果融合器"""
    
    def __init__(self):
        """初始化融合器"""
        pass
    
    def weighted_average_fusion(self, predictions_dict, weights=None):
        """
        加权平均融合多个算法的预测结果
        
        Args:
            predictions_dict: 预测结果字典，格式为 {algorithm_name: predictions}
            weights: 权重字典，格式为 {algorithm_name: weight}，如果为None则使用默认权重
            
        Returns:
            融合后的预测结果列表
        """
        if not predictions_dict:
            return []
        
        # 获取预测组数
        num_predictions = len(list(predictions_dict.values())[0])
        
        # 如果没有指定权重，使用默认权重（基于算法优先级）
        if weights is None:
            weights = self._get_default_weights(predictions_dict.keys())
        
        # 归一化权重
        total_weight = sum(weights.get(algo, 1.0) for algo in predictions_dict.keys())
        if total_weight > 0:
            normalized_weights = {algo: weights.get(algo, 1.0) / total_weight 
                                for algo in predictions_dict.keys()}
        else:
            # 如果权重总和为0，使用均匀权重
            normalized_weights = {algo: 1.0 / len(predictions_dict) 
                                for algo in predictions_dict.keys()}
        
        # 融合预测结果
        fused_predictions = []
        for i in range(num_predictions):
            # 对于每个预测组，融合所有算法的结果
            fused_red = self._fuse_single_prediction(
                {algo: predictions[i]['red'] for algo, predictions in predictions_dict.items()},
                {algo: normalized_weights[algo] for algo in predictions_dict.keys()}
            )
            
            # 蓝球融合（仅对非排列5类型）
            fused_blue = []
            has_blue = any('blue' in predictions[i] and predictions[i]['blue'] 
                          for predictions in predictions_dict.values())
            
            if has_blue:
                fused_blue = self._fuse_blue_prediction(
                    {algo: predictions[i].get('blue', []) for algo, predictions in predictions_dict.items()},
                    {algo: normalized_weights[algo] for algo in predictions_dict.keys()}
                )
            
            # 计算平均置信度
            avg_confidence = np.mean([
                predictions[i].get('confidence', 0.5) 
                for predictions in predictions_dict.values()
            ])
            
            fused_predictions.append({
                'red': fused_red,
                'blue': fused_blue,
                'confidence': avg_confidence,
                'algorithm_weights': normalized_weights
            })
        
        return fused_predictions
    
    def _fuse_single_prediction(self, predictions_by_algo, weights):
        """
        融合单个预测（红球或数字）
        
        Args:
            predictions_by_algo: 各算法的预测结果字典
            weights: 各算法权重字典
            
        Returns:
            融合后的预测结果
        """
        if not predictions_by_algo:
            return []
        
        # 确定预测长度（通常为5个数字）
        pred_length = len(list(predictions_by_algo.values())[0])
        
        # 对每个位置进行加权平均
        fused_result = []
        for pos in range(pred_length):
            weighted_sum = 0.0
            total_weight = 0.0
            
            for algo, prediction in predictions_by_algo.items():
                if pos < len(prediction):
                    weight = weights.get(algo, 1.0 / len(predictions_by_algo))
                    weighted_sum += prediction[pos] * weight
                    total_weight += weight
            
            # 四舍五入到整数
            fused_value = int(round(weighted_sum / total_weight if total_weight > 0 else 0))
            
            # 确保在合理范围内
            if pred_length == 5:  # 排列5
                fused_value = max(0, min(9, fused_value))
            elif pred_length == 6:  # 双色球红球
                fused_value = max(1, min(33, fused_value))
            elif pred_length == 5:  # 大乐透前区
                fused_value = max(1, min(35, fused_value))
            
            fused_result.append(fused_value)
        
        return fused_result
    
    def _fuse_blue_prediction(self, blue_predictions_by_algo, weights):
        """
        融合蓝球预测
        
        Args:
            blue_predictions_by_algo: 各算法的蓝球预测结果字典
            weights: 各算法权重字典
            
        Returns:
            融合后的蓝球预测结果
        """
        if not blue_predictions_by_algo:
            return []
        
        # 收集所有蓝球数字
        all_blue_numbers = []
        all_weights = []
        
        for algo, blue_pred in blue_predictions_by_algo.items():
            if blue_pred:  # 如果有蓝球预测
                weight = weights.get(algo, 1.0 / len(blue_predictions_by_algo))
                all_blue_numbers.extend(blue_pred)
                all_weights.extend([weight] * len(blue_pred))
        
        if not all_blue_numbers:
            return []
        
        # 基于权重选择最可能的蓝球数字
        number_weights = defaultdict(float)
        for num, weight in zip(all_blue_numbers, all_weights):
            number_weights[num] += weight
        
        # 选择权重最高的数字
        sorted_numbers = sorted(number_weights.items(), key=lambda x: x[1], reverse=True)
        
        # 确定蓝球数量（双色球1个，大乐透2个）
        target_count = min(2, len(sorted_numbers))  # 最多2个蓝球
        fused_blue = [num for num, _ in sorted_numbers[:target_count]]
        
        # 排序
        return sorted(fused_blue)
    
    def _get_default_weights(self, algorithm_names):
        """
        获取默认权重（基于算法优先级）
        
        Args:
            algorithm_names: 算法名称列表
            
        Returns:
            权重字典
        """
        # 预定义的算法优先级（数字越小优先级越高）
        priority_map = {
            "lstm_crf": 1,          # 最高优先级
            "enhanced_lstm": 2,     # 高优先级
            "markov_chain": 3,      # 中高优先级
            "gradient_boost": 4,    # 中优先级
            "memory_network": 5,    # 中优先级
            "weighted_expert": 6    # 低优先级
        }
        
        weights = {}
        max_priority = max(priority_map.get(name, 999) for name in algorithm_names)
        
        for name in algorithm_names:
            priority = priority_map.get(name, 999)
            # 权重与优先级成反比（优先级越高权重越大）
            weights[name] = max(0.1, (max_priority - priority + 1) * 0.5)
        
        return weights
    
    def confidence_weighted_fusion(self, predictions_dict):
        """
        基于置信度的加权融合
        
        Args:
            predictions_dict: 预测结果字典
            
        Returns:
            融合后的预测结果列表
        """
        if not predictions_dict:
            return []
        
        # 获取预测组数
        num_predictions = len(list(predictions_dict.values())[0])
        
        fused_predictions = []
        for i in range(num_predictions):
            # 计算每个算法的置信度权重
            confidence_weights = {}
            total_confidence = 0.0
            
            for algo, predictions in predictions_dict.items():
                if i < len(predictions):
                    confidence = predictions[i].get('confidence', 0.5)
                    confidence_weights[algo] = confidence
                    total_confidence += confidence
            
            # 归一化置信度权重
            if total_confidence > 0:
                normalized_weights = {algo: conf / total_confidence 
                                    for algo, conf in confidence_weights.items()}
            else:
                # 如果总置信度为0，使用均匀权重
                normalized_weights = {algo: 1.0 / len(predictions_dict) 
                                    for algo in predictions_dict.keys()}
            
            # 融合预测
            fused_red = self._fuse_single_prediction(
                {algo: predictions[i]['red'] for algo, predictions in predictions_dict.items()},
                normalized_weights
            )
            
            # 蓝球融合
            fused_blue = []
            has_blue = any('blue' in predictions[i] and predictions[i]['blue'] 
                          for predictions in predictions_dict.values())
            
            if has_blue:
                fused_blue = self._fuse_blue_prediction(
                    {algo: predictions[i].get('blue', []) for algo, predictions in predictions_dict.items()},
                    normalized_weights
                )
            
            # 计算平均置信度
            avg_confidence = np.mean([
                predictions[i].get('confidence', 0.5) 
                for predictions in predictions_dict.values()
            ])
            
            fused_predictions.append({
                'red': fused_red,
                'blue': fused_blue,
                'confidence': avg_confidence,
                'algorithm_weights': normalized_weights
            })
        
        return fused_predictions


# 测试代码
if __name__ == "__main__":
    # 创建测试数据
    test_predictions = {
        "lstm_crf": [
            {"red": [3, 5, 4, 5, 6], "blue": [], "confidence": 0.85},
            {"red": [2, 8, 4, 4, 9], "blue": [], "confidence": 0.82}
        ],
        "markov_chain": [
            {"red": [1, 2, 3, 4, 5], "blue": [], "confidence": 0.75},
            {"red": [2, 3, 4, 5, 6], "blue": [], "confidence": 0.72}
        ],
        "memory_network": [
            {"red": [4, 6, 5, 6, 7], "blue": [], "confidence": 0.65},
            {"red": [3, 9, 5, 5, 0], "blue": [], "confidence": 0.62}
        ]
    }
    
    # 创建融合器
    fusion = PredictionFusion()
    
    # 测试加权平均融合
    print("🧪 测试加权平均融合...")
    fused_results = fusion.weighted_average_fusion(test_predictions)
    
    for i, result in enumerate(fused_results):
        print(f"  融合预测 {i+1}: {result['red']} (置信度: {result['confidence']:.3f})")
        print(f"    权重: {result['algorithm_weights']}")
    
    # 测试置信度加权融合
    print("\n🧪 测试置信度加权融合...")
    confidence_fused_results = fusion.confidence_weighted_fusion(test_predictions)
    
    for i, result in enumerate(confidence_fused_results):
        print(f"  置信度融合 {i+1}: {result['red']} (置信度: {result['confidence']:.3f})")
        print(f"    权重: {result['algorithm_weights']}")